#include "MkvgenTestFixture.h"

class AnnexBCpdNalAdapterTest : public MkvgenTestBase {
};

TEST_F(AnnexBCpdNalAdapterTest, nalAdapter_InvalidInput)
{
    PBYTE pCpd = (PBYTE) 100;
    UINT32 cpdSize = 10;
    PBYTE pAdaptedCpd = (PBYTE) 110;
    UINT32 adaptedCpdSize = 100;

    EXPECT_NE(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(NULL, cpdSize, pAdaptedCpd, &adaptedCpdSize));
    EXPECT_NE(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(pCpd, cpdSize, NULL, NULL));
    EXPECT_NE(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(pCpd, cpdSize, pAdaptedCpd, NULL));
    EXPECT_NE(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(pCpd, 0, pAdaptedCpd, &adaptedCpdSize));
    EXPECT_NE(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(pCpd, 1, pAdaptedCpd, &adaptedCpdSize));
    EXPECT_NE(STATUS_SUCCESS,
              adaptH264CpdNalsFromAnnexBToAvcc(pCpd, MIN_H264_ANNEXB_CPD_SIZE - 1, pAdaptedCpd, &adaptedCpdSize));
}

TEST_F(AnnexBCpdNalAdapterTest, nalAdapter_InvalidNoStartCode)
{
    BYTE cpd[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    BYTE adaptedCpd[100];
    UINT32 cpdSize = SIZEOF(cpd);
    UINT32 adaptedSize;

    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd, cpdSize, NULL, &adaptedSize));
    EXPECT_EQ(16, adaptedSize);

    adaptedSize = 100;
    EXPECT_NE(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd, cpdSize, adaptedCpd, &adaptedSize));
    EXPECT_EQ(16, adaptedSize);
}

TEST_F(AnnexBCpdNalAdapterTest, nalAdapter_ValidZerosInStream)
{
    BYTE cpds[][21] = {
            {0, 0, 1, 0x67, 2, 3, 4, 10, 11, 12, 13, 14, 15, 0, 0, 1, 0x68, 3, 4, 5, 6}, // 0
            {0, 0, 0, 1, 0x67, 2, 3, 4, 10, 11, 12, 13, 14, 15, 0, 0, 0, 1, 0x68, 3, 4}, // 1
    };

    UINT32 adaptedSizes[] = {26, 26};
    UINT32 actualAdaptedSizes[] = {26, 24};

    BYTE adaptedCpds[][26] = {
            {1, 2, 3, 4, 0xff, 0xE1, 0, 10, 0x67, 2, 3, 4, 10, 11, 12, 13, 14, 15, 1, 0, 5, 0x68, 3, 4, 5, 6}, // 0
            {1, 2, 3, 4, 0xff, 0xE1, 0, 10, 0x67, 2, 3, 4, 10, 11, 12, 13, 14, 15, 1, 0, 3, 0x68, 3, 4, 0, 0}, // 1
    };

    UINT32 cpdSize;
    BYTE adaptedCpd[1000];
    UINT32 adaptedCpdSize = SIZEOF(adaptedCpd);

    for (UINT32 i = 0; i < SIZEOF(cpds) / SIZEOF(cpds[0]); i++) {
        cpdSize = SIZEOF(cpds[i]);
        EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpds[i], cpdSize, NULL, &adaptedCpdSize));
        EXPECT_EQ(adaptedSizes[i], adaptedCpdSize);
        MEMSET(adaptedCpd, 0x00, SIZEOF(adaptedCpd));
        EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpds[i], cpdSize, adaptedCpd, &adaptedCpdSize));
        EXPECT_EQ(actualAdaptedSizes[i], adaptedCpdSize);
        EXPECT_EQ(0, MEMCMP(adaptedCpds[i], adaptedCpd, adaptedCpdSize)) << "Failed comparison on index: " << i;
    }
}

TEST_F(AnnexBCpdNalAdapterTest, nalAdapter_ValidSpsPps)
{
    BYTE cpd[] = {0x00, 0x00, 0x00, 0x01, 0x67, 0x42, 0x00, 0x1e,
                        0xa9, 0x50, 0x14, 0x07, 0xb4, 0x20, 0x00, 0x00,
                        0x7d, 0x00, 0x00, 0x1d, 0x4c, 0x00, 0x80, 0x00,
                        0x00, 0x00, 0x01, 0x68, 0xce, 0x3c, 0x80};
    UINT32 cpdSize = SIZEOF(cpd);
    BYTE adaptedCpd[1000];
    UINT32 adaptedCpdSize = SIZEOF(adaptedCpd);

    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd, cpdSize, NULL, &adaptedCpdSize));
    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd, cpdSize, adaptedCpd, &adaptedCpdSize));

    BYTE cpd2[] = {0x00, 0x00, 0x00, 0x01, 0x67, 0x64, 0x00, 0x34,
                   0xAC, 0x2B, 0x40, 0x1E, 0x00, 0x78, 0xD8, 0x08,
                   0x80, 0x00, 0x01, 0xF4, 0x00, 0x00, 0xEA, 0x60,
                   0x47, 0xA5, 0x50, 0x00, 0x00, 0x00, 0x01, 0x68,
                   0xEE, 0x3C, 0xB0};
    cpdSize = SIZEOF(cpd2);

    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd2, cpdSize, NULL, &adaptedCpdSize));
    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd2, cpdSize, adaptedCpd, &adaptedCpdSize));
}

TEST_F(AnnexBCpdNalAdapterTest, nalAdapter_FixedUpValidSpsPps)
{
    // This is taken from a real encoder that generates bad NALUs with zeroes at the end.
    // It also has SEI and AUD NALus
    BYTE cpd[] = {0x00, 0x00, 0x00, 0x01, 0x09, 0x10, 0x00, 0x00, 0x00, 0x00, 0x01, 0x67, 0x64, 0x00, 0x1E, 0xAC,
                  0x1B, 0x1A, 0x80, 0xB0, 0x3D, 0xFF, 0xFF, 0x00, 0x28, 0x00, 0x21, 0x6E, 0x0C, 0x0C, 0x0C, 0x80,
                  0x00, 0x01, 0xF4, 0x00, 0x00, 0x75, 0x30, 0x74, 0x30, 0x07, 0xD0, 0x00, 0x01, 0x31, 0x2D, 0x5D,
                  0xE5, 0xC6, 0x86, 0x00, 0xFA, 0x00, 0x00, 0x26, 0x25, 0xAB, 0xBC, 0xB8, 0x50, 0x00, 0x00, 0x00,
                  0x00, 0x01, 0x68, 0xEE, 0x38, 0x30, 0x00, 0x00, 0x00, 0x00, 0x01, 0x06, 0x00, 0x0D, 0xBC, 0xFF,
                  0x87, 0x49, 0xB5, 0x16, 0x3C, 0xFF, 0x87, 0x49, 0xB5, 0x16, 0x40, 0x01, 0x04, 0x00, 0x78, 0x08,
                  0x10, 0x06, 0x01, 0xC4, 0x80,};

    UINT32 cpdSize = SIZEOF(cpd);
    BYTE adaptedCpd[1000];
    UINT32 adaptedCpdSize = SIZEOF(adaptedCpd);

    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd, cpdSize, NULL, &adaptedCpdSize));
    EXPECT_EQ(STATUS_SUCCESS, adaptH264CpdNalsFromAnnexBToAvcc(cpd, cpdSize, adaptedCpd, &adaptedCpdSize));
}

TEST_F(AnnexBCpdNalAdapterTest, nalAdapter_ValidH265) {
    BYTE adaptedCpd[1000];
    UINT32 adaptedCpdSize = SIZEOF(adaptedCpd);

    // These are series of codec private data bits extracted from an
    // Android device for H265 for various resolutions
    BYTE cpd[8][1000] = {
            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x3c, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x3c, 0xa0, 0x18, 0x20, 0x28, 0x71, 0x31, 0x39, 0x6b, 0xb9, 0x32, 0x4b, 0xb9, 0x48, 0x28, 0x10,
                    0x10, 0x17, 0x68, 0x50, 0x94, 0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0xc0, 0xf1, 0x80, 0x04, 0x20},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x3c, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x3c, 0xa0, 0x0a, 0x08, 0x04, 0x07, 0xc4, 0xe5, 0xae, 0xe4, 0xc9, 0x2e, 0xe5, 0x20, 0xa0, 0x40,
                    0x40, 0x5d, 0xa1, 0x42, 0x50, 0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0xc0, 0xf1, 0x80, 0x04, 0x20},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x5a, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x5a, 0xa0, 0x05, 0x02, 0x01, 0xe1, 0x65, 0xae, 0xe4, 0xc9, 0x2e, 0xe5, 0x20, 0xa0, 0x40, 0x40},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x5a, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x5a, 0xa0, 0x06, 0x02, 0x01, 0xc1, 0xf1, 0x39, 0x6b, 0xb9, 0x32, 0x4b, 0xb9, 0x48, 0x28, 0x10,
                    0x10, 0x17, 0x68, 0x50, 0x94, 0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0xc0, 0xf1, 0x80, 0x04, 0x20},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x5a, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x5a, 0xa0, 0x05, 0xc2, 0x01, 0xe1, 0xc4, 0xf9, 0x6b, 0xb9, 0x32, 0x4b, 0xb9, 0x48, 0x28, 0x10,
                    0x10, 0x17, 0x68, 0x50, 0x94, 0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0xc0, 0xf1, 0x80, 0x04, 0x20},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x5d, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x5d, 0xa0, 0x07, 0x82, 0x00, 0xb8, 0x7c, 0x4e, 0x5a, 0xee, 0x4c, 0x92, 0xee, 0x52, 0x0a, 0x04,
                    0x04, 0x05, 0xda, 0x14, 0x25, 0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0xc0, 0xf1, 0x80, 0x04, 0x20},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x5d, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x5d, 0xa0, 0x02, 0x80, 0x80, 0x2e, 0x1f, 0x13, 0x96, 0xbb, 0x93, 0x24, 0xbb, 0x94, 0x82, 0x81,
                    0x01, 0x01, 0x76, 0x85, 0x09, 0x40, 0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0xc0, 0xf1, 0x80, 0x04},

            {0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00,
                    0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x5d, 0xac, 0x59, 0x00, 0x00, 0x00, 0x01, 0x42,
                    0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03, 0x00, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
                    0x5d, 0xa0, 0x02, 0x80, 0x80, 0x3c, 0x16, 0x5a, 0xee, 0x4c, 0x92, 0xee, 0x52, 0x0a, 0x04, 0x04}
    };

    UINT32 cpdSizes[8] = {80, 80, 64, 80, 80, 80, 80, 64};
    UINT32 naluCounts[8] = {3, 3, 2, 3, 3, 3, 3, 2};

    UINT16 i;

    for (i = 0; i < 8; i++) {
        EXPECT_EQ(STATUS_SUCCESS, adaptH265CpdNalsFromAnnexBToHvcc(cpd[i], cpdSizes[i], NULL, &adaptedCpdSize))
                            << "Failed on iteration " << i;
        EXPECT_EQ(STATUS_SUCCESS, adaptH265CpdNalsFromAnnexBToHvcc(cpd[i], cpdSizes[i], adaptedCpd, &adaptedCpdSize))
                            << "Failed on iteration " << i;

        EXPECT_EQ(naluCounts[i], adaptedCpd[HEVC_CPD_HEADER_OVERHEAD - 1])
                            << "Failed on iteration " << i;

        // The avgFrameRate should be 0 = two bytes
        EXPECT_EQ(0, adaptedCpd[HEVC_CPD_HEADER_OVERHEAD - 3])
                            << "Failed on iteration " << i;
        EXPECT_EQ(0, adaptedCpd[HEVC_CPD_HEADER_OVERHEAD - 4])
                            << "Failed on iteration " << i;
    }
}
